import subprocess
import os
import csv
import math

# 创建结果目录
results_dir = '../results_raw'
os.makedirs(results_dir, exist_ok=True)

# 结果文件
output_file = os.path.join(results_dir, 'operator_performance.csv')

# 定义参数组合
op_types = [
    "logit", "boxcox", "onehot", "bucketize", "sigrid_hash",
    "clamp_list", "ngram", "firstx", "fillnull", "cast"
]
device_types = ["arm", "dpa"]
thread_nums_dpa = [2**i for i in range(8)] + [192] + [160] + [12] # 1, 2, 4, ..., 128
thread_nums_cpu_arm = [2**i for i in range(5)] + [12]  # 1, 2, 4, ..., 16, 12
batch_sizes = [2**i for i in range(25)]  # 从1到2^23

# 不支持的操作类型
unsupported_op_types_dpa = {"logit", "onehot", "boxcox", "cast"}

def convert_to_us(value, unit):
    """将时间单位转换为微秒"""
    if unit == 's':
        return value * 1_000_000
    elif unit == 'ms':
        return value * 1_000
    elif unit == 'us':
        return value
    elif unit == 'ns':
        return value / 1000
    else:
        raise ValueError(f"Unknown unit: {unit}")

# 打开文件以写入结果
with open(output_file, 'a', newline='') as csvfile:
    fieldnames = ['op_type', 'device_type', 'thread_num', 'batch_size', 'P50_us']
    writer = csv.DictWriter(csvfile, fieldnames=fieldnames)

    # 如果文件为空，则写入头部
    if csvfile.tell() == 0:
        writer.writeheader()

    for op_type in op_types:
        for device_type in device_types:
            # 根据设备类型选择线程数
            if device_type == "dpa":
                if op_type in unsupported_op_types_dpa:
                    continue
                thread_nums = thread_nums_dpa
            else:
                thread_nums = thread_nums_cpu_arm

            for thread_num in thread_nums:
                for batch_size in batch_sizes:
                    # 提示当前运行的参数组合
                    print(f"Running with op_type={op_type}, device_type={device_type}, thread_num={thread_num}, batch_size={batch_size}")

                    # 构建命令
                    command = [
                        '/home/ubuntu/xeg/BenchBF3/bin_dpu/test_operator',
                        f'--op_type={op_type}',
                        f'--device_type={device_type}',
                        f'--thread_num={thread_num}',
                        f'--batch_size={batch_size}'
                    ]

                    # 执行命令并捕获输出
                    try:
                        result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, encoding='utf-8', errors='ignore')
                    except Exception as e:
                        print(f"Error executing command: {e}")
                        continue

                    stdout = result.stdout
                    lines = stdout.splitlines()
                    for line in lines:
                        if op_type in line:
                            # 提取P50值
                            parts = line.split('|')
                            if len(parts) >= 4:
                                p50_value, p50_unit = parts[2].strip().split()
                                p99_value, p99_unit = parts[3].strip().split()

                                p50_us = convert_to_us(float(p50_value), p50_unit)
                                p99_us = convert_to_us(float(p99_value), p99_unit)

                                writer.writerow({
                                    'op_type': op_type,
                                    'device_type': device_type,
                                    'thread_num': thread_num,
                                    'batch_size': batch_size,
                                    'P50_us': p50_us
                                })

print(f"Results saved to {output_file}")
